package org.farm.framework.elasticsearch.config;

import com.alibaba.fastjson2.JSON;
import com.zozoev.common.enums.LogLevelEnum;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.RandomScoreFunctionBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import javax.annotation.PostConstruct;
import java.util.*;

/**
 * @author Frank Zou
 * @version v0.0.1
 * @date 2024/6/25
 * @description
 */
@Slf4j
@Component
@SuppressWarnings("all")
public class EsRestClient {

    public static Map<String, RestHighLevelClient> clients = new HashMap<>();

    private final EsConfigProperties esConfigProperties;

    private static final RequestOptions COMMON_OPTIONS;

    @Value("${elasticsearch.log-level}")
    private static String logLevel;

    static {
        RequestOptions.Builder builder = RequestOptions.DEFAULT.toBuilder();
        COMMON_OPTIONS = builder.build();
    }

    public EsRestClient(EsConfigProperties esConfigProperties) {
        this.esConfigProperties = esConfigProperties;
    }

    /**
     * 初始化所有的 Elasticsearch 客户端
     */
    @PostConstruct
    public void initialize() {
        List<EsClusterConfig> clusters = esConfigProperties.getClusters();
        for (EsClusterConfig cluster : clusters) {
            log.info("Initializing Elasticsearch client for cluster: name: {}, nodes: {}", cluster.getName(), cluster.getNodes());
            RestHighLevelClient client = initRestHighLevelClient(cluster);
            if (Objects.nonNull(client)) {
                clients.put(cluster.getName(), client);
            } else {
                log.error("Failed to initialize Elasticsearch client for cluster: name: {}, nodes: {}", cluster.getName(), cluster.getNodes());
            }
        }
        log.info("Elasticsearch clients initialized: {}", clients.keySet());
    }

    /**
     * 初始化 Elasticsearch 客户端
     *
     * @param cluster 集群配置
     * @return Elasticsearch 客户端
     */
    private RestHighLevelClient initRestHighLevelClient(EsClusterConfig cluster) {
        String[] nodes = cluster.getNodes().split(",");
        List<HttpHost> httpPostList = new ArrayList<>(nodes.length);
        for (String node : nodes) {
            String[] hostPort = node.split(":");
            if (hostPort.length != 2) {
                log.error("Invalid Elasticsearch node: {}", node);
                return null;
            }
            HttpHost httpPost = new HttpHost(hostPort[0], NumberUtils.toInt(hostPort[1]));
            httpPostList.add(httpPost);
        }
        HttpHost[] httpHosts = new HttpHost[httpPostList.size()];
        httpPostList.toArray(httpHosts);

        // Create credentials provider and set the credentials
        final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(AuthScope.ANY,
                new UsernamePasswordCredentials(cluster.getUsername(), cluster.getPassword()));
        // Build the RestClient with the credentials
        RestClientBuilder builder = RestClient.builder(httpHosts)
                .setHttpClientConfigCallback(httpClientBuilder -> httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider))
                .setRequestConfigCallback(requestConfigBuilder ->
                        requestConfigBuilder.setConnectTimeout(cluster.getConnectTimeout())
                                .setSocketTimeout(cluster.getSocketTimeout()));
        return new RestHighLevelClient(builder);
    }

    /**
     * 获取指定集群的客户端
     *
     * @param clusterName 集群名称
     * @return 集群客户端
     */
    public static RestHighLevelClient getClient(String clusterName) {
        return clients.get(clusterName);
    }


    /**
     * 插入文档
     *
     * @param esIndexInfo  索引信息
     * @param esSourceData 文档数据
     */
    @SneakyThrows
    public static void insertDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData) {
        IndexRequest indexRequest = new IndexRequest(esIndexInfo.getIndexName());
        indexRequest.source(esSourceData.getData());
        indexRequest.id(esSourceData.getDocId());
        IndexResponse response = getClient(esIndexInfo.getClusterName()).index(indexRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
            log.info("Inserted document: {}, response: {}", esSourceData.getDocId(), JSON.toJSONString(response));
        }
    }

    @SneakyThrows
    public static void insertBatchDoc(EsIndexInfo esIndexInfo, List<EsSourceData> esSourceDataList) {
        boolean flag = false;
        BulkRequest bulkRequest = new BulkRequest();
        for (EsSourceData esSourceData : esSourceDataList) {
            if (StringUtils.isBlank(esSourceData.getDocId())) {
                continue;
            }
            IndexRequest indexRequest = new IndexRequest(esIndexInfo.getIndexName());
            indexRequest.source(esSourceData.getData());
            indexRequest.id(esSourceData.getDocId());
            bulkRequest.add(indexRequest);
            flag = true;
        }
        if (flag) {
            BulkResponse bulkResponse = getClient(esIndexInfo.getClusterName()).bulk(bulkRequest, COMMON_OPTIONS);
            if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
                log.info("Inserted documents: {}", esSourceDataList.size());
            }
            if (bulkResponse.hasFailures()) {
                log.error("Failed to insert documents: {}, response: {}", esSourceDataList.size(), JSON.toJSONString(bulkResponse));
            }
        }
    }

    /**
     * 更新文档
     *
     * @param esIndexInfo  集群索引信息
     * @param esSourceData 文档数据
     */
    @SneakyThrows
    public static void updateDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData) {
        UpdateRequest updateRequest = new UpdateRequest();
        updateRequest.index(esIndexInfo.getIndexName())
                .id(esSourceData.getDocId())
                .doc(esSourceData.getData());
        getClient(esIndexInfo.getClusterName()).update(updateRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {

            log.info("Updated document: {}, params: {}", esSourceData.getDocId(), JSON.toJSONString(esSourceData));
        }
    }

    @SneakyThrows
    public static boolean updateBatchDoc(EsIndexInfo esIndexInfo, List<EsSourceData> esSourceDataList) {
        boolean flag = false;
        BulkRequest bulkRequest = new BulkRequest();
        for (EsSourceData esSourceData : esSourceDataList) {
            if (StringUtils.isBlank(esSourceData.getDocId())) {
                continue;
            }
            UpdateRequest updateRequest = new UpdateRequest();
            updateRequest.index(esIndexInfo.getIndexName())
                    .id(esSourceData.getDocId())
                    .doc(esSourceData.getData());
            bulkRequest.add(updateRequest);
            flag = true;
        }
        if (flag) {
            BulkResponse bulkResponse = getClient(esIndexInfo.getClusterName()).bulk(bulkRequest, COMMON_OPTIONS);
            if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {

                log.info("Updated documents: {}, params: {}", esSourceDataList.size(), JSON.toJSONString(esSourceDataList));
            }
            if (bulkResponse.hasFailures()) {
                log.error("Failed to update documents: {}, response: {}", esSourceDataList.size(), JSON.toJSONString(bulkResponse));
                return false;
            }
        }
        return true;
    }

    /**
     * 删除所有文档
     *
     * @param esIndexInfo 集群索引信息
     * @return 是否删除成功
     */
    @SneakyThrows
    public static boolean deleteAllDoc(EsIndexInfo esIndexInfo) {
        DeleteByQueryRequest deleteByQueryRequest = new DeleteByQueryRequest(esIndexInfo.getIndexName());
        deleteByQueryRequest.setQuery(QueryBuilders.matchAllQuery());
        BulkByScrollResponse response = getClient(esIndexInfo.getClusterName()).deleteByQuery(deleteByQueryRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
            log.info("Deleted {} documents from index: {}, response: {}", response.getDeleted(), esIndexInfo.getIndexName(), JSON.toJSONString(response));
        }
        return response.getDeleted() > 0;
    }

    /**
     * 根据 DocID 删除文档
     *
     * @param esIndexInfo 集群索引信息
     * @param docId       Doc id
     * @return 是否删除成功
     */
    @SneakyThrows
    public static boolean deleteDoc(EsIndexInfo esIndexInfo, String docId) {
        DeleteRequest deleteRequest = new DeleteRequest(esIndexInfo.getIndexName());
        deleteRequest.id(docId);
        DeleteResponse delete = getClient(esIndexInfo.getClusterName()).delete(deleteRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
            log.info("Deleted document: {}, response: {}", docId, JSON.toJSONString(delete));
        }
        return true;
    }

    @SneakyThrows
    public static GetResponse getById(EsIndexInfo esIndexInfo, String docId) {
        GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
        getRequest.id(docId);
        GetResponse response = getClient(esIndexInfo.getClusterName()).get(getRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
            log.info("Get document by id: {}, response: {}", docId, JSON.toJSONString(response));
        }
        return response;
    }

    /**
     * 文档是否存在
     *
     * @param esIndexInfo 集群索引信息
     * @param docId       Doc id
     * @return true: 存在; false: 不存在
     */
    @SneakyThrows
    public static boolean isExistsDoc(EsIndexInfo esIndexInfo, String docId) {
        GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
        getRequest.id(docId);
        boolean exists = getClient(esIndexInfo.getClusterName()).exists(getRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
            log.info("Document exists: {}, response: {}", docId, exists);
        }
        return exists;
    }

    /**
     * 根据 DocID 获取文档
     *
     * @param esIndexInfo 集群索引信息
     * @param docId       Doc id
     * @return 文档数据
     */
    @SneakyThrows
    public static Map<String, Object> getDocById(EsIndexInfo esIndexInfo, String docId) {
        GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
        getRequest.id(docId);
        GetResponse response = getClient(esIndexInfo.getClusterName()).get(getRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
            log.info("Get document by id: {}, response: {}", docId, JSON.toJSONString(response));
        }
        return response.getSource();
    }


    /**
     * 根据 DocID 获取文档
     *
     * @param esIndexInfo 集群索引信息
     * @param docId       Doc id
     * @param includes    查询结果需要包含什么字段
     * @param excludes    查询结果需要排除什么字段
     * @return 文档数据
     */
    @SneakyThrows
    public static Map<String, Object> getDocFieldById(EsIndexInfo esIndexInfo, String docId, String[] includes, String[] excludes) {
        GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
        getRequest.id(docId);
        FetchSourceContext fetchSourceContext = new FetchSourceContext(true, includes, excludes);
        getRequest.fetchSourceContext(fetchSourceContext);
        GetResponse response = getClient(esIndexInfo.getClusterName()).get(getRequest, COMMON_OPTIONS);
        if (LogLevelEnum.DEBUG.getLevel().equalsIgnoreCase(logLevel)) {
            log.info("Get document field by id: {}, response: {}", docId, JSON.toJSONString(response));
        }
        return response.getSource();
    }

    /**
     * 搜索文档
     *
     * @param esIndexInfo     集群索引信息
     * @param esSearchRequest 搜索请求条件
     * @return 搜索结果
     */
    @SneakyThrows
    public static SearchResponse searchWithTermQuery(EsIndexInfo esIndexInfo, EsSearchRequest esSearchRequest) {
        BoolQueryBuilder queryBuilder = esSearchRequest.getQueryBuilder();
        AggregationBuilder aggregationBuilder = esSearchRequest.getAggregationBuilder();
        String[] fields = esSearchRequest.getFields();
        Integer from = esSearchRequest.getFrom();
        Integer size = esSearchRequest.getSize();
        Boolean needScroll = esSearchRequest.getNeedScroll();
        Long minute = esSearchRequest.getMinute();
        List<FieldSortBuilder> fieldSortBuilders = esSearchRequest.getFieldSortBuilders();
        HighlightBuilder highlightBuilder = esSearchRequest.getHighlightBuilder();
        Boolean isRandomScore = esSearchRequest.getIsRandomScore();
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.trackTotalHits(true);
        searchSourceBuilder.query(queryBuilder);
        searchSourceBuilder.fetchSource(fields, null);
        // 根据 From 判断是否需要分页
        if (Objects.nonNull(from)) {
            searchSourceBuilder.from(from).size(size);
        }
        if (Objects.nonNull(highlightBuilder)) {
            searchSourceBuilder.highlighter(highlightBuilder);
        }

        if (!CollectionUtils.isEmpty(fieldSortBuilders)) {
            // 如果排序中不是空的，则添加排序
            fieldSortBuilders.forEach(searchSourceBuilder::sort);
        }else{
            if (isRandomScore) {
                // 添加随机排序
                FunctionScoreQueryBuilder functionScoreQueryBuilder = QueryBuilders.functionScoreQuery(queryBuilder, new RandomScoreFunctionBuilder());
                searchSourceBuilder.query(functionScoreQueryBuilder);
            }
        }

        // 添加聚合查询
        if (Objects.nonNull(aggregationBuilder)) {
            searchSourceBuilder.aggregation(aggregationBuilder);
        }
        SearchRequest searchRequest = new SearchRequest();
        searchRequest.searchType(SearchType.DEFAULT);
        searchRequest.indices(esIndexInfo.getIndexName());
        searchRequest.source(searchSourceBuilder);
        if (needScroll) {
            searchRequest.scroll(new Scroll(TimeValue.timeValueMinutes(minute)));
        }
        return getClient(esIndexInfo.getClusterName()).search(searchRequest, COMMON_OPTIONS);
    }


}
